#!/usr/bin/python
# -*- coding: utf-8 -*-
# Description: PE处理器
# Create: 2008-6-1
# Author: MK2[fengmk2@gmail.com]
import os
import time
from datetime import datetime
from fileauth.client.filesliceapp.runtaskbase import RunTaskBase
from fileauth.public.utility import common as util
from fileauth.public.utility import peutil
os.environ["DJANGO_SETTINGS_MODULE"] = "fileauth.web.settings" 
from django.db import connection
from fileauth.web.infocollectapp import infosaver 
from fileauth.client.infocollectapp.infocollect import InfoCollect

class PEProcesser(RunTaskBase):
    """样本PE处理者"""
    @property
    def tablename(self):
        return 'fileslicemgr_pewaiting88'
    
    @property
    def type(self):
        return 8088
    
    @property
    def limit_num(self):
        return 100
    
    SELECT_PEWAITING_LIST = 'SELECT md5, local_path from %s WHERE id >= %d AND id <= %d'
    def task(self, cursor):
        cursor.execute(self.__class__.SELECT_PEWAITING_LIST % \
                       (self.tablename, self.cur_id, self.max_id))
        rs = cursor.fetchall()
        self._collect_infos(rs, cursor)

    def log(self, *args):
        now = datetime.now()
        logdir = './log/%s' % now.strftime('%Y/%m/')
        util.ensureDirExist(logdir)
        logpath = os.path.join(logdir, now.strftime('%d.log'))
        self.locker.lock()
        log = open(logpath, 'a')
        log.write('|'.join(map(str, args)))
        log.write(os.linesep)
        log.close()
        self.locker.unlock()
        
    INSERT_BAD_PE = "INSERT INTO fileslicemgr_badpe(md5) VALUES(%s)"
    INSERT_XML_LOST = "INSERT INTO fileslicemgr_canotcreatexml(md5) VALUES(%s)"
    DELETE_PEWAITING_LIST = 'DELETE FROM %s WHERE id>=%d AND id<=%d'
    def _collect_infos(self, md5_path_list, cursor):
        """md5_path_list: [[md5, local_path], ...]"""
        xml_lost = []
        xml_success = []
        self.info('Creating the PE xml file...')
        infos, md5_xmlpath_list, fail_md5s, badpes, \
            hasverion_count, hasdigital_count, has_vd_count, \
            notpe_count, totalsize, usetime = InfoCollect.getInfos(md5_path_list)
        uploadresults, xmlfiles_size, uploadtimes = \
                        InfoCollect.upload_infofiles(md5_xmlpath_list)
        for key, val in uploadresults.iteritems():
            if val == -1:
                xml_lost.append(key)
            else:
                xml_success.append(key)
        infosaver.savemany(infos, connection)
        self.info('PE: %d, has versioninfo: %d, digital: %d, success upload: %d, fail upload: %d' \
                    % (len(infos), hasverion_count, hasdigital_count, 
                       len(xml_success), len(xml_lost)))
        self.info('Bad pe: %d, not pe: %d' % (len(badpes), notpe_count))
        
        #(4)保存未能创建样本的md5到二库
        cursor.executemany(self.INSERT_XML_LOST, xml_lost)
        cursor.executemany(self.INSERT_BAD_PE, badpes)
        #pecount, hasverion_count, digital, successupload, failupload, 
        #badpe, notpe, xmlsize, xmlavgsize, uptime, xmlavguptime, 
        #totalsize, peavgsize, usetime, peavgtime
        self.log(*(len(infos), hasverion_count, hasdigital_count, 
                  len(xml_success), len(xml_lost),
                  len(badpes), notpe_count,
                  xmlfiles_size, xmlfiles_size / len(md5_xmlpath_list),
                  uploadtimes, uploadtimes * 1.0 / len(md5_xmlpath_list),
                  totalsize, totalsize * 1.0 / len(md5_path_list), 
                  usetime, usetime * 1.0 / len(md5_path_list), )) 
        #(5)删除本地样本及未能删除的xml文件
#        cursor.execute(self.DELETE_PEWAITING_LIST % (self.tablename, self.cur_id, self.max_id)) #删除记录
#        for info in fileinfos:
#            try:
#                if os.path.isfile(info[1]): 
#                    os.remove(info[1])
#            except Exception, e:
#                self.info('Could not delete the file: %s, error: %s' % (info[1], e))

import sys 
if __name__=='__main__':
    task_id = 0
    if len(sys.argv) > 1:
        task_id = int(sys.argv[1])
    PEProcesser(task_id).run()